package cn.iocoder.yudao.module.aiposter.service.vec;

import cn.iocoder.yudao.framework.common.pojo.PageResult;
import cn.iocoder.yudao.framework.common.util.json.JsonUtils;
import cn.iocoder.yudao.framework.common.util.object.BeanUtils;
import cn.iocoder.yudao.module.aiposter.controller.app.template.vo.AppTemplatePageReqVO;
import cn.iocoder.yudao.module.aiposter.controller.app.vec.vo.VecTemplateMixSimilaritySearchReqVO;
import cn.iocoder.yudao.module.aiposter.controller.app.vec.vo.VecTemplateSimilaritySearchReqVO;
import cn.iocoder.yudao.module.aiposter.controller.app.vec.vo.VecTemplateSimilaritySearchRespVO;
import cn.iocoder.yudao.module.aiposter.dal.dataobject.template.TemplateDO;
import cn.iocoder.yudao.module.aiposter.remote.VecClient;
import cn.iocoder.yudao.module.aiposter.remote.vo.PosterFilterBO;
import cn.iocoder.yudao.module.aiposter.remote.vo.PosterVecMeta;
import cn.iocoder.yudao.module.aiposter.remote.vo.VecDBItemReq;
import cn.iocoder.yudao.module.aiposter.remote.vo.VecDBResp;
import cn.iocoder.yudao.module.aiposter.service.template.AppTemplateService;
import cn.iocoder.yudao.module.aiposter.utils.CommonUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;

@Service
@Slf4j
public class VecServiceImpl implements VecService {

    private static Double SIMILARITY_SEARCH_THRESHOLD = 0.79d;

    @Resource
    private VecClient vecClient;
    @Resource
    private AppTemplateService appTemplateService;

    @Override
    public Boolean resetIndex() {
        // 删除所有海报索引
        vecClient.cleanAll();
        // 查询所有模版 构建插入信息插入数据
        AppTemplatePageReqVO pageReq = new AppTemplatePageReqVO();

        Integer pageNum = 0;
        boolean isFinish = false;
        while (!isFinish) {
            pageNum++;
            pageReq.setPageNo(pageNum);
            PageResult<TemplateDO> templatePage = appTemplateService.getTemplatePage(pageReq);
            List<TemplateDO> list = templatePage.getList();
            if (null != list && list.size() > 0) {
                List<VecDBItemReq> addVecList = new ArrayList<>();
                list.forEach(item -> {
                    VecDBItemReq reqItem = new VecDBItemReq();
                    reqItem.setTemplateId(item.getId());
                    reqItem.setBizName(item.getBizName());
                    reqItem.setKeyword(item.getKeyword());
                    reqItem.setIsAI(appTemplateService.judgeIsAI(item.getData()));
                    reqItem.setIsSegmentation(appTemplateService.judgeIsSegmentation(item.getData()));
                    addVecList.add(reqItem);
                });
                // 分批插入向量库
                vecClient.add(addVecList);
            } else {
                isFinish = true;
            }
        }
        return true;
    }

    @Override
    public List<VecTemplateSimilaritySearchRespVO> similaritySearch(VecTemplateSimilaritySearchReqVO reqVO) {
        List<String> requestItem = new ArrayList<String>() {{
            add(reqVO.getKeyword());
        }};
        Integer numResult = 4;
        PosterFilterBO posterFilterBO = PosterFilterBO.builder()
                .bizName(reqVO.getBizName())
                .isAI(reqVO.getIsAI())
                .isSegmentation(reqVO.getIsSegmentation())
                .build();
        List<VecDBResp.SimilaritySearchItem> isResultList = vecClient.similaritySearch(requestItem, numResult, posterFilterBO);

        List<VecTemplateSimilaritySearchRespVO> result = new ArrayList<>();
        for (int i = 0; i < isResultList.get(0).getSearch().size(); i++) {
            PosterVecMeta metadata = getVecMetaWithScoreThreshold(isResultList.get(0).getSearch().get(i));
            if (metadata != null) {
                VecTemplateSimilaritySearchRespVO resultItem = BeanUtils.toBean(metadata, VecTemplateSimilaritySearchRespVO.class);
                result.add(resultItem);
            }
        }
        return result;
    }

    /**
     * 相似度查询，返回3个抠图合成模版 1个AI合成模版
     *
     * 为了满足每次随机获取，根据相似度匹配分别召回0个和10个最相似的模版，然后随机挑选3个和1个
     * @param reqVO 请求参数
     * @return 相似的海报模版列表
     */
    @Override
    public List<VecTemplateSimilaritySearchRespVO> similaritySearchMix(VecTemplateMixSimilaritySearchReqVO reqVO) {
        // 最大召回AI结果
        int maxAIResult = 10;
        // 最大召回Segmentation结果
        int maxSegmentationResult = 30;
        // 生成AI召回随机数个数
        int generateAI = 1;
        // 生成Seg召回随机数个数
        int generateSeg = 3;
        List<String> requestItem = new ArrayList<String>() {{
            add(reqVO.getKeyword());
        }};
        PosterFilterBO posterFilterBO = PosterFilterBO.builder()
                .bizName(reqVO.getBizName())
                .isAI(true)
                .isSegmentation(false)
                .build();
        List<VecDBResp.SimilaritySearchItem> isAIResultList = vecClient.similaritySearch(requestItem, maxAIResult, posterFilterBO);
        posterFilterBO.setIsSegmentation(true);
        posterFilterBO.setIsAI(false);
        List<VecDBResp.SimilaritySearchItem> isSegResultList = vecClient.similaritySearch(requestItem, maxSegmentationResult, posterFilterBO);

        // 过滤阈值
        List<PosterVecMeta> aiMetaList = new ArrayList<>();
        for(List<Object> item : isAIResultList.get(0).getSearch()){
            PosterVecMeta metadata = getVecMetaWithScoreThreshold(item);
            if(null != metadata){
                aiMetaList.add(metadata);
            }
        }
        List<PosterVecMeta> segMetaList = new ArrayList<>();
        for(List<Object> item : isSegResultList.get(0).getSearch()){
            PosterVecMeta metadata = getVecMetaWithScoreThreshold(item);
            if(null != metadata){
                segMetaList.add(metadata);
            }
        }

        // 3个Seg结果
        List<VecTemplateSimilaritySearchRespVO> result = new ArrayList<>();
        int segRealResult = Math.min(segMetaList.size(), maxSegmentationResult);
        List<Integer> randomResult = CommonUtils.generateRandomNumbers(generateSeg, 0 ,segRealResult);
        for (Integer randomIndex : randomResult) {
            PosterVecMeta metadata = segMetaList.get(randomIndex);
            VecTemplateSimilaritySearchRespVO resultItem = BeanUtils.toBean(metadata, VecTemplateSimilaritySearchRespVO.class);
            result.add(resultItem);
        }
        // 1个AI结果
        int aiRealResult = Math.min(aiMetaList.size(), maxAIResult);
        randomResult = CommonUtils.generateRandomNumbers(generateAI, 0 ,aiRealResult);
        for (Integer randomIndex : randomResult) {
            PosterVecMeta metadata = aiMetaList.get(randomIndex);
            VecTemplateSimilaritySearchRespVO resultItem = BeanUtils.toBean(metadata, VecTemplateSimilaritySearchRespVO.class);
            result.add(resultItem);
        }
        return result;
    }

    /**
     * 根据向量召回的结果SimilaritySearchItem.search中的某一个条目的Meta信息
     * @param searchItemInfo 关键词向量召回的结果列表
     * @return 关键词下结果条目的Meta信息
     */
    private PosterVecMeta getVecMetaWithScoreThreshold(List<Object> searchItemInfo) {
        VecDBResp.MatchInfo matchInfo = JsonUtils.parseObject(JsonUtils.toJsonString(searchItemInfo.get(0)),VecDBResp.MatchInfo.class);
        PosterVecMeta metadata = matchInfo.getMetadata();
        Double sameScore = (Double) searchItemInfo.get(1);
        if (sameScore > SIMILARITY_SEARCH_THRESHOLD) {
            // 底相似度内容不渲染
            return null;
        }
        return metadata;
    }
}
